#!/usr/bin/env python3
"""检查 CUDA 和 ONNX Runtime 配置"""

print("=" * 60)
print("检查 CUDA 和 ONNX Runtime 配置")
print("=" * 60)

# 1. 检查 PyTorch CUDA
print("\n1. 检查 PyTorch CUDA 支持:")
try:
	import torch

	print(f"   PyTorch 版本: {torch.__version__}")
	print(f"   CUDA 可用: {torch.cuda.is_available()}")
	if torch.cuda.is_available():
		print(f"   CUDA 版本: {torch.version.cuda}")
		print(f"   GPU 数量: {torch.cuda.device_count()}")
		print(f"   GPU 名称: {torch.cuda.get_device_name(0)}")
except ImportError:
	print("   ❌ PyTorch 未安装")
except Exception as e:
	print(f"   ❌ 错误: {e}")

# 2. 检查 ONNX Runtime
print("\n2. 检查 ONNX Runtime:")
try:
	import onnxruntime

	print(f"   ONNX Runtime 版本: {onnxruntime.__version__}")
	print(f"   可用提供商: {onnxruntime.get_available_providers()}")

	if 'CUDAExecutionProvider' in onnxruntime.get_available_providers():
		print("   ✅ CUDA 执行提供商可用")
	else:
		print("   ❌ CUDA 执行提供商不可用")
		print("\n   原因可能是:")
		print("   1. onnxruntime-gpu 版本与 CUDA 版本不匹配")
		print("   2. CUDA 或 cuDNN 未正确安装")
		print("   3. 环境变量未正确设置")
except ImportError:
	print("   ❌ ONNX Runtime 未安装")
except Exception as e:
	print(f"   ❌ 错误: {e}")

# 3. 检查已安装的包
print("\n3. 检查已安装的相关包:")
try:
	import subprocess

	result = subprocess.run(['pip', 'list'], capture_output=True, text=True)
	lines = result.stdout.split('\n')

	print("   相关包:")
	for line in lines:
		if any(pkg in line.lower() for pkg in ['onnx', 'torch', 'cuda']):
			print(f"   - {line}")
except Exception as e:
	print(f"   ❌ 错误: {e}")

# 4. 检查 CUDA 环境变量
print("\n4. 检查 CUDA 环境变量:")
import os

cuda_vars = ['CUDA_PATH', 'CUDA_HOME', 'CUDA_VISIBLE_DEVICES']
for var in cuda_vars:
	value = os.environ.get(var, '未设置')
	print(f"   {var}: {value}")

print("\n" + "=" * 60)
print("检查完成")
print("=" * 60)

# 5. 给出建议
print("\n建议:")
print("如果 CUDA 不可用，请尝试:")
print("1. 卸载当前版本: pip uninstall onnxruntime onnxruntime-gpu")
print("2. 重新安装 GPU 版本: pip install onnxruntime-gpu")
print("3. 或使用 FaceFusion 官方安装器: python install.py --onnxruntime cuda")
print("\n注意: onnxruntime-gpu 需要:")
print("- CUDA 11.x 或 12.x")
print("- cuDNN 8.x")
print("- 对应的 NVIDIA 驱动")

